#!/usr/bin/env python3
"""
ComfyUI 依赖管理工具 (跨平台版本)
功能：
1. 安装/更新依赖库
2. 卸载所有第三方依赖库（自动备份）
3. 查看已安装的依赖库
4. 备份当前依赖库
5. 还原依赖库
6. 智能依赖解析和冲突处理
"""

import os
import sys
import subprocess
import platform
import shutil
from datetime import datetime
from pathlib import Path
from typing import List, Dict, Optional, Tuple
import logging
from .env_setup import ENV_CONFIG  # 修复为相对导入

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='[%(levelname)s] %(message)s',
    handlers=[
        logging.StreamHandler(sys.stdout),
        logging.FileHandler('dependency_manager.log', encoding='utf-8')
    ]
)
logger = logging.getLogger(__name__)

class DependencyManager:
    """依赖管理插件类"""
    
    # 插件排序
    order = 1
    
    def __init__(self):
        """初始化依赖管理器"""
        # 使用环境配置中的路径
        self.env_config = ENV_CONFIG
        self.workspace_path = ENV_CONFIG['PROJECT_ROOT']
        
        # 直接使用当前运行的Python解释器
        self.python_path = sys.executable
        
        self.backup_dir = self.workspace_path / "dependency_backups"
        self.backup_dir.mkdir(exist_ok=True)
    
    def start(self) -> bool:
        """插件启动方法"""
        print("🔧 依赖管理工具")
        print("=" * 40)
        
        # 显示基本信息
        print(f"🐍 Python 版本: {sys.version}")
        print(f"⚙️ Python 解释器路径: {self.python_path}")
        print(f"📂 当前工作路径: {self.workspace_path}")
        
        # 运行主程序
        self.run()
        return True
    
    @staticmethod
    def Launch(self):
        self.run()

    def _run_command(self, cmd: List[str], check: bool = True) -> Tuple[int, str, str]:
        """运行命令并返回结果"""
        try:
            # 根据操作系统选择合适的编码
            if platform.system() == "Windows":
                encoding = 'gbk'  # Windows中文环境使用GBK编码
            else:
                encoding = 'utf-8'  # Unix系统使用UTF-8编码
            
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                encoding=encoding,
                errors='replace'  # 遇到无法解码的字符时替换为占位符
            )
            if check and result.returncode != 0:
                logger.error(f"命令失败: {' '.join(cmd)}")
                logger.error(f"错误输出: {result.stderr}")
            return result.returncode, result.stdout, result.stderr
        except Exception as e:
            logger.error(f"执行命令失败: {e}")
            return -1, "", str(e)
    
    def _check_admin_privileges(self) -> bool:
        """检查是否有管理员权限"""
        if platform.system() == "Windows":
            try:
                return subprocess.run(
                    ["net", "session"], 
                    capture_output=True, 
                    check=False
                ).returncode == 0
            except:
                return False
        else:
            # Unix系统检查是否为root
            return os.geteuid() == 0
    
    def _create_backup_filename(self, backup_name: str) -> str:
        """创建备份文件名"""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        return f"{backup_name}_{timestamp}.txt"
    
    def _get_backup_files_sorted(self) -> List[Path]:
        """获取按时间排序的备份文件列表（最近10个）"""
        backup_files = list(self.backup_dir.glob("*.txt"))
        # 按修改时间排序，最新的在前
        backup_files.sort(key=lambda x: x.stat().st_mtime, reverse=True)
        return backup_files[:10]  # 只返回最近10个
    
    def show_menu(self):
        """显示主菜单"""
        print("\n" + "="*50)
        print("           ComfyUI 依赖管理工具")
        print("="*50)
        print()
        print("请选择操作：")
        print("1. 安装/更新依赖库")
        print("2. 卸载所有第三方依赖库")
        print("3. 查看已安装的依赖库")
        print("4. 备份当前依赖库")
        print("5. 还原依赖库")
        print("6. 检查权限")
        print("7. 退出")
        print()
    
    def install_dependencies(self):
        """安装/更新依赖库"""
        print("\n🚀 开始安装/更新依赖库...\n")
        
        # 检查并安装基础工具
        tools = ["pip-tools", "pipdeptree", "piptree","packaging"]
        for tool in tools:
            print(f"检查 {tool}...")
            code, _, _ = self._run_command([self.python_path, "-m", "pip", "show", tool], check=False)
            if code != 0:
                print(f"安装 {tool}...")
                self._run_command([self.python_path, "-m", "pip", "install", tool])
            else:
                print(f"✅ {tool} 已安装")
        
        # 直接import depend并调用其main函数
        try:
            import sys
            import importlib
            
            # 使用环境配置中的路径
            paths = [
                str(self.env_config['COMFYUI_DIR']),
                str(self.env_config['COMFYUI_CUSTOM_NODES']),
                str(self.env_config['CONFIG_CUSTOM_NODES'])
            ]
            print(f"扫描路径: {paths}")
            
            # 直接调用depend模块（环境已配置好）
            from . import depend
            sys.argv = ["depend.py"] + paths
            depend.main()
            print("✅ 依赖安装完成")
        except Exception as e:
            print("❌ 依赖安装失败（import调用）")
            print(f"错误: {e}")
    
    def uninstall_dependencies(self):
        """卸载所有第三方依赖库"""
        print("\n⚠️  警告：即将卸载所有第三方依赖库！")
        print("\n此操作将：")
        print("- 自动备份当前所有已安装的包列表")
        print("- 卸载所有第三方依赖库（保留 pip, setuptools, wheel）")
        print("- 清理 Python 环境")
        print()
        
        confirm = input("确认要卸载所有第三方依赖库吗？(y/N): ").strip().lower()
        if confirm != 'y':
            print("操作已取消")
            return
        
        # 自动备份
        print("\n📁 自动备份当前包列表...")
        backup_file = self._create_backup_filename("auto_backup")
        backup_path = self.backup_dir / backup_file
        
        code, stdout, _ = self._run_command([self.python_path, "-m", "pip", "freeze"])
        if code == 0:
            with open(backup_path, 'w', encoding='utf-8') as f:
                f.write(stdout)
            print(f"✅ 备份已保存到: {backup_path}")
        else:
            print("❌ 备份失败")
            return
        
        # 获取第三方包列表
        print("\n🔍 获取第三方包列表...")
        code, stdout, _ = self._run_command([self.python_path, "-m", "pip", "freeze"])
        if code != 0:
            print("❌ 获取包列表失败")
            return
        
        # 过滤第三方包
        third_party_packages = []
        for line in stdout.strip().split('\n'):
            line = line.strip()
            if line and not any(skip in line.lower() for skip in ['pip', 'setuptools', 'wheel']):
                # 跳过注释行和Git可编辑安装
                if not line.startswith('#') and not line.startswith('-e '):
                    third_party_packages.append(line)
        
        if not third_party_packages:
            print("✅ 没有找到需要卸载的第三方包")
            return
        
        # 显示将要卸载的包
        print(f"\n📋 将要卸载 {len(third_party_packages)} 个包：")
        for i, pkg in enumerate(third_party_packages, 1):
            print(f"  {i}. {pkg}")
        
        confirm2 = input(f"\n确认卸载以上 {len(third_party_packages)} 个包吗？(y/N): ").strip().lower()
        if confirm2 != 'y':
            print("操作已取消")
            return
        
        # 卸载包
        print("\n🗑️  正在卸载第三方包...")
        
        # 批量卸载：将所有包名合并成一个命令
        pkg_names = []
        for pkg in third_party_packages:
            # 处理Git可编辑安装的包
            if pkg.startswith('-e '):
                # 对于 -e path/to/package 格式，提取路径中的包名
                path_part = pkg[3:].strip()
                if '\\' in path_part or '/' in path_part:
                    # 从路径中提取包名（通常是最后一个目录名）
                    pkg_name = path_part.replace('\\', '/').split('/')[-1]
                else:
                    pkg_name = path_part
            elif pkg.startswith('#'):
                # 跳过注释行
                continue
            elif '==' in pkg:
                # 标准格式：package==version
                pkg_name = pkg.split('==')[0]
            elif '>=' in pkg:
                # 版本范围格式：package>=version
                pkg_name = pkg.split('>=')[0]
            elif '<=' in pkg:
                # 版本范围格式：package<=version
                pkg_name = pkg.split('<=')[0]
            elif '!=' in pkg:
                # 版本排除格式：package!=version
                pkg_name = pkg.split('!=')[0]
            elif '~=' in pkg:
                # 兼容版本格式：package~=version
                pkg_name = pkg.split('~=')[0]
            else:
                # 没有版本信息的包名
                pkg_name = pkg.strip()
            
            # 清理包名（去除可能的空格和特殊字符）
            pkg_name = pkg_name.strip()
            if pkg_name and not pkg_name.startswith('#'):
                pkg_names.append(pkg_name)
        
        print(f"🚀 开始批量卸载 {len(pkg_names)} 个包...")
        
        # 分批处理，每批20个包，避免命令行过长
        batch_size = 20
        success_count = 0
        total_batches = (len(pkg_names) + batch_size - 1) // batch_size
        
        for batch_num in range(total_batches):
            start_idx = batch_num * batch_size
            end_idx = min(start_idx + batch_size, len(pkg_names))
            batch_packages = pkg_names[start_idx:end_idx]
            
            print(f"📦 批次 {batch_num + 1}/{total_batches}: 卸载 {len(batch_packages)} 个包...")
            
            # 构建批量卸载命令
            uninstall_cmd = [self.python_path, "-m", "pip", "uninstall", "-y"] + batch_packages
            
            code, stdout, stderr = self._run_command(uninstall_cmd, check=False)
            
            if code == 0:
                print(f"✅ 批次 {batch_num + 1} 卸载成功")
                success_count += len(batch_packages)
            else:
                print(f"⚠️  批次 {batch_num + 1} 部分失败，尝试逐个卸载...")
                print(f"   错误信息: {stderr.strip() if stderr else '未知错误'}")
                
                # 分析失败原因
                if "command line too long" in stderr.lower():
                    print("   💡 原因: 命令行过长，建议减少批次大小")
                elif "permission" in stderr.lower():
                    print("   💡 原因: 权限不足，某些包需要管理员权限")
                elif "dependency" in stderr.lower():
                    print("   💡 原因: 包之间存在依赖关系，需要逐个处理")
                
                # 如果批量失败，回退到逐个卸载
                batch_success = 0
                for pkg_name in batch_packages:
                    print(f"   🔧 逐个卸载: {pkg_name}")
                    code, _, _ = self._run_command([self.python_path, "-m", "pip", "uninstall", "-y", pkg_name], check=False)
                    if code == 0:
                        success_count += 1
                        batch_success += 1
                    else:
                        print(f"      ❌ {pkg_name} 卸载失败")
                
                print(f"   📊 批次 {batch_num + 1} 逐个卸载结果: {batch_success}/{len(batch_packages)} 成功")
        
        print(f"\n✅ 卸载完成: {success_count}/{len(third_party_packages)} 个包成功卸载")
        print(f"📁 备份文件: {backup_path}")
        print("\n如需恢复，可以运行选项 5 进行还原")
    
    def list_dependencies(self):
        """查看已安装的依赖库"""
        print("\n📋 已安装的依赖库列表：\n")
        
        code, stdout, _ = self._run_command([self.python_path, "-m", "pip", "list"])
        if code == 0:
            print(stdout)
        else:
            print("❌ 获取包列表失败")
    
    def backup_dependencies(self):
        """备份当前依赖库"""
        print("\n📁 备份当前依赖库...\n")
        
        # 显示最近10条备份记录
        recent_backups = self._get_backup_files_sorted()
        if recent_backups:
            print("📋 最近10条备份记录：")
            for i, backup_file in enumerate(recent_backups, 1):
                # 获取文件修改时间
                mtime = datetime.fromtimestamp(backup_file.stat().st_mtime)
                time_str = mtime.strftime("%Y-%m-%d %H:%M:%S")
                print(f"  {i}. {backup_file.name} ({time_str})")
            print()
        
        backup_name = input("请输入备份名称（可选，直接回车使用默认名称）: ").strip()
        if not backup_name:
            backup_name = "manual_backup"
        
        backup_file = self._create_backup_filename(backup_name)
        backup_path = self.backup_dir / backup_file
        
        print(f"📁 备份当前包列表到: {backup_path}")
        
        code, stdout, _ = self._run_command([self.python_path, "-m", "pip", "freeze"])
        if code == 0:
            with open(backup_path, 'w', encoding='utf-8') as f:
                f.write(stdout)
            print("✅ 备份创建成功")
            print(f"📁 备份文件: {backup_path}")
        else:
            print("❌ 备份创建失败")
    
    def restore_dependencies(self):
        """还原依赖库"""
        print("\n🔄 还原依赖库...\n")
        
        # 获取最近10条备份记录
        backup_files = self._get_backup_files_sorted()
        
        if not backup_files:
            print("❌ 未找到任何备份文件")
            print("请先创建备份或检查备份文件是否存在")
            return
        
        # 显示可用的备份文件（最近10个）
        print("📋 最近10条备份记录：")
        for i, backup_file in enumerate(backup_files, 1):
            # 获取文件修改时间
            mtime = datetime.fromtimestamp(backup_file.stat().st_mtime)
            time_str = mtime.strftime("%Y-%m-%d %H:%M:%S")
            print(f"{i}. {backup_file.name} ({time_str})")
        
        try:
            choice = input(f"\n请选择要还原的备份文件 [1-{len(backup_files)}]: ").strip()
            if not choice:
                return
            
            choice_num = int(choice)
            if choice_num < 1 or choice_num > len(backup_files):
                print("❌ 无效选择")
                return
            
            selected_backup = backup_files[choice_num - 1]
            
            print(f"\n⚠️  警告：即将还原依赖库！")
            print(f"此操作将覆盖当前已安装的包")
            print(f"备份文件: {selected_backup.name}")
            print()
            
            confirm = input("确认要还原吗？(y/N): ").strip().lower()
            if confirm != 'y':
                print("操作已取消")
                return
            
            print("\n🔄 正在还原依赖库...")
            code, stdout, stderr = self._run_command([self.python_path, "-m", "pip", "install", "-r", str(selected_backup)])
            
            if code == 0:
                print("✅ 依赖库还原成功")
            else:
                print("❌ 依赖库还原失败")
                print(f"错误: {stderr}")
                
        except ValueError:
            print("❌ 无效输入")
        except Exception as e:
            print(f"❌ 操作失败: {e}")
    
    def check_privileges(self):
        """检查权限"""
        print("\n🔐 权限检查...\n")
        
        is_admin = self._check_admin_privileges()
        if is_admin:
            print("✅ 已获得管理员权限")
        else:
            print("❌ 未获得管理员权限")
            print("某些操作可能需要管理员权限")
        
        print(f"🐍 Python路径: {self.python_path}")
        print(f"📂 工作目录: {self.workspace_path}")
        print(f"💾 备份目录: {self.backup_dir}")
    
    def run(self):
        """运行主程序"""
        print(f"🐍 Python 版本: {sys.version}")
        print(f"⚙️ Python 解释器路径: {self.python_path}")
        print(f"📂 当前工作路径: {self.workspace_path}")
        
        while True:
            self.show_menu()
            
            try:
                choice = input("请输入数字 [1-7]: ").strip()
                
                if choice == "1":
                    self.install_dependencies()
                elif choice == "2":
                    self.uninstall_dependencies()
                elif choice == "3":
                    self.list_dependencies()
                elif choice == "4":
                    self.backup_dependencies()
                elif choice == "5":
                    self.restore_dependencies()
                elif choice == "6":
                    self.check_privileges()
                elif choice == "7":
                    print("\n退出程序")
                    break
                else:
                    print("❌ 无效选择，请重新输入")
                
                input("\n按回车键继续...")
                
            except KeyboardInterrupt:
                print("\n\n退出程序")
                break
            except Exception as e:
                print(f"\n❌ 发生错误: {e}")
                input("按回车键继续...")

def main():
    """主函数"""
    # 获取脚本所在目录
    script_dir = Path(__file__).parent
    workspace_path = script_dir
    
    # 创建依赖管理器
    manager = DependencyManager()
    
    # 运行主程序
    manager.start()

if __name__ == "__main__":
    main() 